package com.xj.driver.mqtt.subscriber;


import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.eclipse.paho.client.mqttv3.*;
import org.eclipse.paho.client.mqttv3.persist.MemoryPersistence;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLong;

/**
 * EMQX压力测试工具 - 每分钟5000条数据
 */
public class EmqxPressureTest {
    private static final Logger log = LoggerFactory.getLogger(EmqxPressureTest.class);

    // EMQX连接配置
    private static final String BROKER = "tcp://11.100.100.201:1883"; // 修改为您的EMQX地址
    private static final String CLIENT_ID = "pressure_test_client";
    private static final String TOPIC = "zhny_iot_data/cn_cust_datas/up/biz/A:43010001_01";

    // 测试配置
    private static final int TOTAL_MESSAGES = 100; // 总消息数
    private static final int THREAD_COUNT = 10;     // 并发线程数
    private static final int DURATION_MINUTES = 1;  // 测试时长(分钟)

    private final AtomicLong sentCount = new AtomicLong(0);
    private final AtomicLong errorCount = new AtomicLong(0);
    private final Random random = new Random();
    private final SimpleDateFormat timeFormat = new SimpleDateFormat("yyyyMMddHHmmssSSS");

    public static void main(String[] args) {
        new EmqxPressureTest().startTest();
    }

    public void startTest() {
        log.info("🚀 开始EMQX压力测试");
        log.info("目标: {}分钟内发送 {} 条消息", DURATION_MINUTES, TOTAL_MESSAGES);
        log.info("主题: {}", TOPIC);
        log.info("EMQX地址: {}", BROKER);

        ExecutorService executor = Executors.newFixedThreadPool(THREAD_COUNT);
        CountDownLatch latch = new CountDownLatch(THREAD_COUNT);

        long startTime = System.currentTimeMillis();
        long endTime = startTime + DURATION_MINUTES * 60 * 1000;

        // 创建多个生产者线程
        for (int i = 0; i < THREAD_COUNT; i++) {
            final int threadId = i;
            executor.submit(() -> {
                try {
                    MqttClient client = createMqttClient(threadId);
                    produceMessages(client, threadId, endTime);
                    client.disconnect();
                    client.close();
                } catch (Exception e) {
                    log.error("线程 {} 执行失败: {}", threadId, e.getMessage());
                } finally {
                    latch.countDown();
                }
            });
        }

        // 启动监控线程
        startMonitorThread(startTime, endTime);

        try {
            latch.await();
            executor.shutdown();

            long totalTime = System.currentTimeMillis() - startTime;
            log.info("✅ 测试完成!");
            log.info("总发送消息: {}", sentCount.get());
            log.info("失败消息: {}", errorCount.get());
            log.info("总耗时: {} ms", totalTime);
            log.info("平均TPS: {}/s", (sentCount.get() * 1000.0 / totalTime));

        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    private MqttClient createMqttClient(int threadId) throws MqttException {
        String clientId = CLIENT_ID + "_" + threadId + "_" + System.currentTimeMillis();
        MqttClient client = new MqttClient(BROKER, clientId, new MemoryPersistence());

        MqttConnectOptions options = new MqttConnectOptions();
        options.setCleanSession(true);
        options.setConnectionTimeout(10);
        options.setKeepAliveInterval(60);
        options.setAutomaticReconnect(true);

        // 设置回调
        client.setCallback(new MqttCallback() {
            @Override
            public void connectionLost(Throwable cause) {
                log.warn("连接断开: {}", cause.getMessage());
            }

            @Override
            public void messageArrived(String topic, MqttMessage message) {
                // 不需要处理接收消息
            }

            @Override
            public void deliveryComplete(IMqttDeliveryToken token) {
                // 消息发布完成回调
            }
        });

        client.connect(options);
        log.info("线程 {} MQTT客户端连接成功: {}", threadId, clientId);
        return client;
    }

    private void produceMessages(MqttClient client, int threadId, long endTime) {
        int messagesPerThread = TOTAL_MESSAGES / THREAD_COUNT;
        int extraMessages = TOTAL_MESSAGES % THREAD_COUNT;

        // 最后一个线程处理剩余的消息
        if (threadId == THREAD_COUNT - 1) {
            messagesPerThread += extraMessages;
        }

        log.info("线程 {} 负责发送 {} 条消息", threadId, messagesPerThread);

        for (int i = 0; i < messagesPerThread; i++) {
            // 检查是否超时
            if (System.currentTimeMillis() > endTime) {
                log.warn("线程 {} 因超时停止发送，已发送: {}/{}", threadId, i, messagesPerThread);
                break;
            }

            try {
                String message = generateMessageData(threadId, i);
                MqttMessage mqttMessage = new MqttMessage(message.getBytes());
                mqttMessage.setQos(1); // QoS 1 确保消息送达
                mqttMessage.setRetained(false);

                client.publish(TOPIC, mqttMessage);
                sentCount.incrementAndGet();

                // 控制发送速率，避免瞬间压力过大
                if (sentCount.get() % 100 == 0) {
                    Thread.sleep(10); // 每100条消息稍微休息一下
                }

            } catch (Exception e) {
                errorCount.incrementAndGet();
                log.error("线程 {} 发送消息失败: {}", threadId, e.getMessage());

                // 短暂休息后重试
                try {
                    Thread.sleep(100);
                } catch (InterruptedException ex) {
                    Thread.currentThread().interrupt();
                    break;
                }
            }
        }

        log.info("线程 {} 完成消息发送: {}/{}", threadId, messagesPerThread, messagesPerThread);
    }

    /**
     * 生成测试消息数据
     */
    private String generateMessageData(int threadId, int messageIndex) {
        JSONObject message = new JSONObject();
        JSONArray devices = new JSONArray();
        JSONObject device = new JSONObject();

        // 生成设备ID
        String deviceId = "device_" + threadId + "_" + messageIndex;
        device.put("deviceId", deviceId);

        JSONArray services = new JSONArray();
        JSONObject service = new JSONObject();
        JSONObject data = new JSONObject();
        JSONObject dData = new JSONObject();

        // 生成随机数据
        dData.put("4010001", random.nextInt(100));        // 0-99的随机整数
        dData.put("4010002", -273.2 + random.nextDouble() * 100); // -273.2 到 -173.2 的随机浮点数
        dData.put("4010003", random.nextInt(50));         // 0-49的随机整数
        dData.put("4010004", random.nextInt(200));        // 0-199的随机整数
        dData.put("4010005", random.nextInt(1000));       // 0-999的随机整数

        String currentTime = timeFormat.format(new Date());
        String eventTime = currentTime.substring(0, 8) + "T" +
                currentTime.substring(8, 10) +
                currentTime.substring(10, 12) +
                currentTime.substring(12, 14) + "Z";

        data.put("station_no", "A:43010001_03");
        data.put("time", currentTime);
        data.put("d_id", "A:43010001_03");
        data.put("d_data", dData);
        data.put("eventTime", eventTime);
        data.put("serviceId", "");

        service.put("data", data);
        services.add(service);
        device.put("services", services);
        devices.add(device);
        message.put("devices", devices);

        return message.toJSONString();
    }

    /**
     * 监控线程，定期输出进度
     */
    private void startMonitorThread(long startTime, long endTime) {
        Thread monitorThread = new Thread(() -> {
            while (System.currentTimeMillis() < endTime) {
                try {
                    Thread.sleep(5000); // 每5秒输出一次进度

                    long currentSent = sentCount.get();
                    long currentErrors = errorCount.get();
                    long elapsed = System.currentTimeMillis() - startTime;
                    double progress = (double) currentSent / TOTAL_MESSAGES * 100;
                    double tps = (currentSent * 1000.0) / elapsed;

                    log.info("📊 进度: {}/{} ({:.2f}%), TPS: {:.2f}/s, 错误: {}",
                            currentSent, TOTAL_MESSAGES, progress, tps, currentErrors);

                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    break;
                }
            }
        });

        monitorThread.setDaemon(true);
        monitorThread.start();
    }
}